#include "txr_algo_dlm_v5_o.h"

using namespace vino_v6;

namespace txr_algo_dlm
{
    TxrAlgoDlmV5_o::TxrAlgoDlmV5_o()
        :m_p_detect(TK_NULL_PTR)
    {
    }

    TxrAlgoDlmV5_o::~TxrAlgoDlmV5_o()
    {
        Release();
    }

    /**
     * 使用提供的配置加载模型。
     * @param cfg 包含模型文件路径的配置。
     * @return 成功返回true，否则返回false。
     */
    tk_bool TxrAlgoDlmV5_o::Load(st_dlm_cfg & cfg)
    {
        m_p_detect = new Yolo_VINO();
        const std::string path_weight = std::string(cfg.file_path_name) + MODEL_FILE_SUFFIX;

        std::cout << "Loading weight: " << path_weight << std::endl;

        const tk_bool rt = m_p_detect->Init(path_weight.c_str());
        if (!rt)
        {
            std::cerr << "Failed to load weight." << std::endl;
            Release();
        }
        std::cout << "Loading Success" << std::endl;
        return rt;
    }

    /**
     * 释放资源。
     */
    void TxrAlgoDlmV5_o::Release()
    {
        SAFE_DELETE(m_p_detect);
    }

    /**
     * 获取批处理大小。
     * @return 批处理大小，未初始化则返回0。
     */
    tk_int32 TxrAlgoDlmV5_o::BatchSize()
    {
        if (m_p_detect)
        {
            return m_p_detect->BatchSize();
        }
        return 0;
    }

    /**
     * 将边界框转换为st_box类型。
     * @param p_box 指向边界框数组的指针。
     * @param box_num 边界框的数量。
     * @param result 存储转换结果的结构体。
     * @param row 图像高度。
     * @param col 图像宽度。
     */
    void BoxTrans(Yolo_VINO::DetRes_VINO * p_box, tk_int32 box_num, st_result_boxes & result, tk_int32 row, tk_int32 col)
    {
        result.vaild_num = box_num;
        for (tk_int32 i = 0; i < box_num; ++i)
        {
            st_box & box = result.boxes[i];
            box.x = p_box[i].x - p_box[i].w / 2;
            box.y = p_box[i].y - p_box[i].h / 2;
            box.w = p_box[i].w;
            box.h = p_box[i].h;
            box.prob = p_box[i].prob;
            box.obj_id = p_box[i].classes;

            box.x = qBound(0, box.x, (col - 1, box.x));
            box.y = qBound(0, box.y, (row - 1, box.y));
            box.w = qBound(0, box.w, (box.w, col - 1 - box.x));
            box.h = qBound(0, box.h, (box.h, row - 1 - box.y));
        }
    }

    /**
     * 执行目标检测。
     * @param p_unit 指向输入检测单元的指针。
     * @param size 输入单元的数量。
     */
    void TxrAlgoDlmV5_o::Detect(st_detect_unit * p_unit, tk_int32 size) {
        if (size <= 0 || p_unit == nullptr) return;

        auto* dlm_data_array = new st_dlm_data[size];

        for (int i = 0; i < size; ++i) {
            const st_detect_unit & unit = p_unit[i];
            st_dlm_data & dlm_data = dlm_data_array[i];

            dlm_data.img.rgb8 = cv::Mat(unit.img.h, unit.img.w, CV_8UC3);

            for (int row = 0; row < unit.img.h; ++row) {
                for (int col = 0; col < unit.img.w; ++col) {
                    const int idx = row * unit.img.w + col;
                    auto &pixel = dlm_data.img.rgb8.at<cv::Vec3b>(row, col);
                    pixel[0] = unit.img.pp_rgb8[2][idx];
                    pixel[1] = unit.img.pp_rgb8[1][idx];
                    pixel[2] = unit.img.pp_rgb8[0][idx];
                }
            }
        }

        std::vector<std::vector<Yolo_VINO::DetRes_VINO>> vv_results;

        m_p_detect->Detect(dlm_data_array, size, vv_results);

        delete[] dlm_data_array;

        for (int i = 0; i < size; ++i)
        {
            const tk_int32 result_box_num = __min(vv_results[i].size(), C_SHARP_MAX_OBJECTS);
            st_img_rgb & rgb = p_unit[i].img;
            BoxTrans(vv_results[i].data(), result_box_num, p_unit[i].result, rgb.h, rgb.w);
        }
    }

    /**
     * 创建并返回一个TxrAlgoDlmV5_o对象。
     * @return 指向新对象的指针。
     */
    extern "C" TXR_ALGO_DLM_V5_O_EXPORT TxrAlgoDlmImpl* CreateDlmObj()
    {
        return new TxrAlgoDlmV5_o();
    }

    /**
     * 安全删除一个TxrAlgoDlmImpl对象。
     * @param p 指向要删除的对象的指针。
     */
    extern "C" TXR_ALGO_DLM_V5_O_EXPORT void DeleteDlmObj(TxrAlgoDlmImpl* p)
    {
        SAFE_DELETE(p);
    }

}